#ifndef RING_H
#define RING_H
#include <cstdint>
#include <vector>
#include <gmpxx.h>
#include "ring_poly.h"
#include "modular_reduction.h"
using namespace std;
/* *
 * @class Ring
 * @brief 表示环（Ring）数据结构的类，包含用于多项式操作和NTT的相关参数与方法。
 *
 * `Ring` 类提供了加密算法中常用的环操作，包括加法、乘法、模运算以及快速傅里叶变换（NTT）相关操作。
 */
class Ring {
public:
    uint64_t N;
    long logN;
    std::vector<uint64_t> Modulus;
    std::vector<uint64_t> Mask;
    bool allowsNTT;
    mpz_class *modulusBigint;
    std::vector<std::vector<uint64_t>> bredParams;
    std::vector<std::vector<uint64_t>> rescaleParams;
    std::vector<uint64_t> psiMont;
    std::vector<uint64_t> psiInvMont;
    std::vector<std::vector<uint64_t>> nttPsi;
    std::vector<std::vector<uint64_t>> nttPsiInv;
    std::vector<uint64_t> nttNInv;
    uint64_t NthRoot;

    // NTT params
    long M;
    uint64_t **qRootPows;
    uint64_t **qRootPowsInv;
    uint64_t **qRootPowsShoup;    // shoup参数
    uint64_t **qRootPowsInvShoup; // shoup参数
    uint64_t *nInvModqShoup;      // shoup参数
    uint64_t *nInvModq;
    uint64_t *qRoots;
    uint64_t *qRootsInv;

    uint64_t *qdVec;
    __uint128_t *qrVec;

    uint64_t *twoqVec;
    uint64_t *qrVecHigh;     // Barrett recution
    uint64_t *qrVecLow;      // Barrett recution
    uint64_t *nScaleInvModq; // [i]

    void NTT(Poly *p1, Poly *p2); // ring_ntt
    void NTTLvl(int level, Poly *p1, Poly *p2);
    void InvNTT(Poly *p1, Poly *p2);
    void InvNTTLvl(uint64_t level, Poly *p1, Poly *p2);

    void iNTTAndEqual(vector<uint64_t> &a, long index);
    void iINTTAndEqual(vector<uint64_t> &a, long index);

    Ring(uint64_t NParam, const std::vector<uint64_t> &Moduli);
    // 拷贝构造函数
    Ring(const Ring &) = delete;
    Ring &operator = (const Ring &) = delete;
    // 移动构造函数
    Ring(Ring &&other) = default;
    Ring &operator = (Ring &&other) = default;
    ~Ring();
    Poly *NewPoly() const;
    Poly *NewPoly(int nbModulus);
    void Permute(const Poly *polIn, uint64_t gen, Poly *polOut);

    void Copy(const Poly *p0, Poly *p1);    // ring_automorphism.
    void Add(Poly *p1, Poly *p2, Poly *p3); // ring_operation
    void AddLvl(int level, Poly *p1, Poly *p2, Poly *p3);
    void AddNoMod(Poly *p1, Poly *p2, Poly *p3);
    void AddNoModLvl(int level, Poly *p1, Poly *p2, Poly *p3);
    void Sub(Poly *p1, Poly *p2, Poly *p3);
    void SubLvl(int level, Poly *p1, Poly *p2, Poly *p3);
    void SubNoMod(Poly *p1, Poly *p2, Poly *p3);
    void SubNoModLvl(int level, Poly *p1, Poly *p2, Poly *p3);
    void Neg(Poly *p1, Poly *p2);
    void NegLvl(int level, Poly *p1, Poly *p2);
    void Reduce(Poly *p1, Poly *p2);
    void ReduceLvl(uint64_t level, Poly *p1, Poly *p2);

    // for NTT
    void SetNTTParameters(uint64_t NNttParam, const std::vector<uint64_t> &modulusParam);
    void SetBarrettParameters(uint64_t NBarrettParam, const std::vector<uint64_t> &modulusParam);

    void MulCoeffsBarrett(Poly *p1, Poly *p2, Poly *p3);
    void MulCoeffsBarrettLvl(uint64_t level, Poly *p1, Poly *p2, Poly *p3);
    void MulCoeffsBarrettAndAdd(Poly *p1, Poly *p2, Poly *p3);
    void MulCoeffsBarrettAndAddNoMod(Poly *p1, Poly *p2, Poly *p3);
    void MulCoeffsBarrettAndSubLvl(int level, Poly *p1, Poly *p2, Poly *p3);
    void MulCoeffsBarrettConstant(Poly *p1, Poly *p2, Poly *p3);

    void MulCoeffs(Poly *p1, Poly *p2, Poly *p3);
    void MulCoeffsLvl(int level, Poly *p1, Poly *p2, Poly *p3);
    void AddScalarBigint(Poly *p1, mpz_class scalar, Poly *p2);
    void SubScalarBigint(Poly *p1, mpz_class scalar, Poly *p2);
    void MulScalar(Poly *p1, uint64_t scalar, Poly *p2);
    void MulScalarLvl(int level, Poly *p1, uint64_t scalar, Poly *p2);

    void MulScalarBarrett(Poly *p1, uint64_t scalar, Poly *p2);
    void MulScalarLvlBarrett(int level, Poly *p1, uint64_t scalar, Poly *p2);

    void SetParameters(uint64_t nParam, const std::vector<uint64_t> &modulusParam);

    int MinLevelBinary(Poly *p1, Poly *p2);
    int MinLevelTernary(Poly *p1, Poly *p2, Poly *p3);
    Poly *NewPolyLvl(int level);
    std::vector<uint64_t> PermuteNTTIndex(uint64_t galEl);

    void PermuteNTTWithIndexLvl(int level, Poly *polIn, std::vector<uint64_t> &index, Poly *polOut);
    void MulScalarBigintLvl(int level, Poly *p1, mpz_class *scalar, Poly *p2);

    void MulCoeffsBarrettConstantLvl(int level, Poly *p1, Poly *p2, Poly *p3);
    void MulCoeffsBarrettConstantAndAddNoModLvl(int level, Poly *p1, Poly *p2, Poly *p3);
};
/* *
 * @class RingQP
 * @brief 表示两个环（`Ring Q` 和 `Ring P`）的组合类，通常用于加密算法中的多项式操作。
 *
 * `RingQP` 类提供了对两个不同环（Q 和 P）上的多项式操作，包括基转换、NTT操作和多项式乘法等。
 */
class RingQP {
public:
    Ring *ringQ;
    Ring *ringP;

    RingQP(Ring *ringQ, Ring *ringP);
    void ExtendBasisSmallNormAndCenter(Poly *polyInQ, int levelP, Poly *polyOutQ, Poly *polyOutP);
    PolyQP *NewPoly();
    PolyQP *NewPolyLvl(int levelQ, int levelP);
    void NTTLazyLvl(int levelQ, int levelP, PolyQP *p, PolyQP *pOut);
    void MulCoeffsBarrettAndSubLvl(int levelQ, int levelP, PolyQP *p1, PolyQP *p2, PolyQP *p3);
    void MulCoeffsBarrettConstantLvl(int levelQ, int levelP, PolyQP *p1, PolyQP *p2, PolyQP *p3);
    void MulCoeffsBarrettConstantAndAddNoModLvl(int levelQ, int levelP, PolyQP *p1, PolyQP *p2, PolyQP *p3);
};

#endif // RING_H